library(survival)
library(survminer)
library(forestplot)

opt <- NULL

opt$outdir <- "PanelScreen"




#cox regresssion model
plotHR <- function(text,res,prefix){
  text[1,] <- gsub("univ_|multi_","",colnames(text))
  pdf(paste0(opt$outdir,'/',prefix,'_forestplot.pdf'),width=10, height=5)
 # pdf("forestplot.pdf",width=960, height=(1+dim(res)[1])*50)
  colinfo <- lapply(1:nrow(res),function(x){ifelse(x %% 2==1,
                                                   return(gpar(col="gray",lineend="butt",columns=c(2:6))),return(gpar(col='gray',lineend="butt",columns=c(2:6))))})
  names(colinfo) <- 1:nrow(res)+1
  forestplot(as.matrix(text),mean=c(NA,log2(res$mean)),lower = c(NA,log2(res$lower)),upper = c(NA,log2(res$upper)),
             graph.pos = 4,graphwidth = unit(60,'mm'),fn.ci_norm="fpDrawDiamondCI",lineheight = unit(12,'mm'),line.margin = unit(5,'mm'),
             colgap = unit(4,'mm'),zero = 0,cex=0.9,col=fpColors(box="#00A896", lines="#02C39A", zero = "black"),
             txt_gp = fpTxtGp(label=gpar(cex=1.25),ticks=gpar(cex=1.1),xlab=gpar(cex=1.2),title = gpar(cex=1.2)),
             lwd.ci=1.5,ci.vertices = T,title="log2 Hazard Ratio",boxsize = 0.3,hrzl_lines = colinfo)
  dev.off()
}

univ_analysis <- function(features,da){
  univ_formulas <- sapply(features,function(x) as.formula(paste('Surv(Time, Dead)~', paste0('`',x,'`'))))
  univ_models <- lapply( univ_formulas, function(x){coxph(x, data = da)})
  univ_results <- lapply(univ_models,
                         function(x){
                           x <- summary(x)
                           p.value<-signif(x$wald["pvalue"])
                           beta<-signif(x$coef[1]);#coeficient beta
                           HR <-signif(x$coef[2]);#exp(beta)
                           HR.confint.lower <- signif(x$conf.int[,"lower .95"])
                           HR.confint.upper <- signif(x$conf.int[,"upper .95"])
                           CI <- paste0(HR.confint.lower, "-", HR.confint.upper)
                           res<-c(beta, HR, CI, p.value)
                           names(res)<-c("univ_beta", "univ_HR", "univ_95% CI for HR", "univ_p.value")
                           return(res)
                           #return(exp(cbind(coef(x),confint(x))))
                         })
  res <- as.data.frame(t(as.data.frame(univ_results, check.names = FALSE)))
  res$univ_beta <- sapply(res$univ_beta,function(x) as.numeric(as.character(x)))
  res$univ_HR <- sapply(res$univ_HR,function(x) as.numeric(as.character(x)))
  res$`univ_95% CI for HR` <- sapply(res$`univ_95% CI for HR`,as.character)
  res$`univ_p.value` <- sapply(res$`univ_p.value`,function(x) as.numeric(as.character(x)))
  return(res)
}

multi_cox <- function(sig_features,data){
  sig_features_formula <- unlist(sapply(sig_features,function(x){paste0('`',x,'`')}))
  multi_formula <- as.formula(paste('Surv(Time, Dead) ~',paste(sig_features_formula,collapse='+')))
  multi_models <- coxph(multi_formula, data = data)
  suma <- summary(multi_models)
  rownames(suma$coef) <- sig_features
  rownames(suma$conf.int) <- sig_features
  multi_res <- sapply(sig_features,function(x){
    p.value <- signif(suma$coef[x,'Pr(>|z|)'])
    beta <- signif(suma$coef[x,'coef'])
    HR <- signif(suma$coef[x,'exp(coef)'])
    HR.confint.lower <- signif(suma$conf.int[x,"lower .95"])
    HR.confint.upper <- signif(suma$conf.int[x,"upper .95"])
    CI <- paste0(HR.confint.lower, "-", HR.confint.upper)
    res<-c(beta, HR, CI, p.value)
    names(res)<-c("multi_beta", "multi_HR", "multi_95% CI for HR", "multi_p.value")
    return(res)
  })
  multi_res <- as.data.frame(t(multi_res))
  multi_res$multi_beta <- sapply(multi_res$multi_beta,function(x) as.numeric(as.character(x)))
  multi_res$multi_HR <- sapply(multi_res$multi_HR,function(x) as.numeric(as.character(x)))
  multi_res$`multi_95% CI for HR` <- sapply(multi_res$`multi_95% CI for HR`,as.character)
  multi_res$`multi_p.value` <- sapply(multi_res$`multi_p.value`,function(x) as.numeric(as.character(x)))
  return(multi_res)
}

forest_plot <- function(res,prefix){
  result <- cbind(rownames(res),res)
  colnames(result)[1] <- "features"
  write.table(result,paste0(opt$outdir,'/',prefix,'_cox_univ_res.txt'),sep="\t",quote=F,row.names=F)
  res$univ_beta <- signif(res$univ_beta,2)
  res$univ_HR <- signif(res$univ_HR,2)
  res$`univ_95% CI for HR` <- sapply(res$`univ_95% CI for HR`,function(x){paste0(signif(as.numeric(unlist(strsplit(x,'-'))),2),collapse = "-")})
  res$univ_p.value <- signif(res$univ_p.value,4)
  res$mean <- as.numeric(as.character(res$univ_HR))
  res$lower <- as.numeric(as.character(sapply(res$`univ_95% CI for HR`,function(x){unlist(strsplit(as.character(x),'-'))[1]})))
  res$upper <- as.numeric(as.character(sapply(res$`univ_95% CI for HR`,function(x){unlist(strsplit(as.character(x),'-'))[2]})))
  text <- rbind(colnames(res)[1:4],as.matrix(res[1:4]))
  rownames(text)[1] <- 'features'
  text <- cbind(rownames(text),text)
  plotHR(text,res,prefix)
}

forest_plot_multi <- function(res,prefix){
  result <- cbind(rownames(res),res)
  colnames(result)[1] <- "features"
  write.table(result,paste0(opt$outdir,'/',prefix,'_cox_multi_res.txt'),sep="\t",quote=F,row.names=F)
  res$multi_beta <- signif(res$multi_beta,2)
  res$multi_HR <- signif(res$multi_HR,2)
  res$`multi_95% CI for HR` <- sapply(res$`multi_95% CI for HR`,function(x){paste0(signif(as.numeric(unlist(strsplit(x,'-'))),2),collapse = "-")})
  res$multi_p.value <- signif(res$multi_p.value,4)  
  res$mean <- as.numeric(as.character(res$multi_HR))
  res$lower <- as.numeric(as.character(sapply(res$`multi_95% CI for HR`,function(x){unlist(strsplit(as.character(x),'-'))[1]})))
  res$upper <- as.numeric(as.character(sapply(res$`multi_95% CI for HR`,function(x){unlist(strsplit(as.character(x),'-'))[2]})))
  text <- rbind(colnames(res)[1:4],as.matrix(res[1:4]))
  rownames(text)[1] <- 'features'
  text <- cbind(rownames(text),text)
  plotHR(text,res,prefix)
}







features <- c("age", "stage","histological_type",
"MetStatus",'GNAQ',	'GNA11',	'EIF1AX',	'SF3B1',"gene_riskscore","cluster_riskscore","DL_riskscore")

#write.csv(new_SHH,"test_data.csv")

Dat <-read.csv("TCGA_allclinical.csv",header = T,row.names = 1,check.names = F)
Dat <-Dat [,9:25]
colnames(Dat)[1:2]<-c("Time","Dead")
Dat$gender <- as.numeric(factor(Dat$gender))
Dat$MetStatus <- ifelse(Dat$MetStatus=="Metastatic",1,0)
Dat$stage <- as.numeric(factor(Dat$stage))
Dat$histological_type <- as.numeric(factor(Dat$histological_type))
#Dat$SCNA_Cluster <- as.numeric(factor(Dat$SCNA_Cluster))
#Dat$chromosome.3.status <- as.numeric(factor(Dat$chromosome.3.status))

clinical_univ_res <- univ_analysis(features,Dat)

clinical_multiv_res<-multi_cox(features,Dat)

forest_plot(clinical_univ_res,"uni_clinical_info")

forest_plot_multi(clinical_multiv_res,"multi_clinical_info")

forest_plot(univ_sig_res,"gene_forest")


